import openai
import requests

from flask import request, jsonify, Response,stream_with_context
from flask import Flask, request
import json
from args import *
from flask_cors import cross_origin
import time
app = Flask(__name__)
app.debug = False
order_of_model = 1

client = openai.Client(
    api_key="cannot be empty",
    base_url=f"http://{xinference_url}:{xinference_port}/v1"
)
def chat_model(prompt, order, model_uid=model_uid, max_tokens=max_tokens, temperature=temperature):

    response = client.chat.completions.create(
        model=f'{model_uid}_{order}',
        messages=[
            {
                "content": prompt,
                "role": "user",
            }
        ],
        max_tokens=max_tokens,
        temperature=temperature,
        stream=True
    )
    return response

def get_lower_len(data):
    return min(len(data), max_input_len)

def wrap_(d):
    if not d:
        return ""
    return d


def wrap(d):
    title = d.get("title", "")
    content = d.get("content", "")
    title = wrap_(title)
    content = wrap_(content)
    return title + " " + content


def get_prompt(source, query):
    source = [f"新闻 {order + 1}：{wrap(i)}" for order, i in enumerate(source)]
    data = '|'.join(source)
    data = data[:get_lower_len(data)]
    data = "说中文：" + data + " " + query
    return data

def wrap_data(content: str):
    ret = {
        "choices": [
            {
                "delta": {
                    "content": content
                }
            }
        ]
    }
    return ret


def event_stream(req):
    global order_of_model
    order_of_model = (order_of_model + 1) % num_of_model + 1
    ids = get(req, "ids", "")
    if not ids:
        yield f'data:{json.dumps(wrap_data("您没有勾选数据，以下将针对你的提问做出回答：<br/>"), ensure_ascii=False)}\n\n'
        data = ""
    else:
        data = requests.get(f'{ids_request_url}?ids={ids}').json()['data']
    query = get(req, "query", "")
    prompt = get_prompt(data, query)
    for r in chat_model(f"{prompt}", order_of_model):
        choices = r.choices
        content = choices[0].delta.content
        yield f'data:{json.dumps(wrap_data(content), ensure_ascii=False)}\n\n'
    yield "data:[DONE]\n\n"

def get(dict_, k,  default):
    if k in dict_:
        return default if dict_[k] is None else dict_[k]
    return default

@app.route('/model/chat_with_ids', methods=['POST'])
@cross_origin()
def post_http():
    req = request.get_json()
    return Response(stream_with_context(event_stream(req)), mimetype="text/event-stream")
    # 返回JSON数据。


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=port)
    # 这里指定了地址和端口号。也可以不指定地址填0.0.0.0那么就会使用本机地址ip
